package _0_ratelimit_demo

import (
	"net/http"
	"time"

	ratelimit2 "github.com/juju/ratelimit"

	ratelimit1 "go.uber.org/ratelimit" //漏桶

	"github.com/gin-gonic/gin"
)

func pintHandler(ctx *gin.Context) {
	ctx.String(http.StatusOK, "pong")
}

func heiHandler(ctx *gin.Context) {
	ctx.String(http.StatusOK, "ha")
}

// 基于漏桶的限流中间件1
func rateLimit1() func(ctx *gin.Context) {
	//生成一个限流器
	rl := ratelimit1.New(100)
	return func(ctx *gin.Context) {
		if rl.Take().Sub(time.Now()) > 0 { //取水滴，返回的是下一滴水的时间
			//time.Sleep(rl.Take().Sub(time.Now())) //需要等这么长时间下一滴水才会滴下来
			ctx.String(http.StatusOK, "rate time...")
			ctx.Abort()
			return
		}
	}
}

// 基于漏桶的限流中间件2
func rateLimit2(fillInterval time.Duration, cap int64) func(ctx *gin.Context) {
	rl := ratelimit2.NewBucket(fillInterval, cap)
	return func(ctx *gin.Context) {
		//rl.Take()          //这一次可以欠账，后面的人就会多一次等一次的时间
		//rl.TakeAvailable() //有令牌的时候才会取出令牌
		if rl.TakeAvailable(1) == 1 { //此次取到令牌
			ctx.Next()
			return
		}
		ctx.String(http.StatusOK, "rate limit...")
		ctx.Abort()
	}
}

func main() {
	r := gin.Default()

	r.GET("/ping", rateLimit1(), pintHandler)
	r.GET("/hei", rateLimit2(time.Second, 10), heiHandler)

	r.Run()
}
